import json
import pandas as pd
import matplotlib.pyplot as plt
from gluonts.dataset.pandas import PandasDataset
from gluonts.dataset.split import split
from huggingface_hub import hf_hub_download
from uni2ts.eval_util.plot import plot_single
from uni2ts.model.moirai import MoiraiForecast, MoiraiModule
import torch

# 示例数据
data = {
    "symbol": ["EURUSD", "EURUSD", "EURUSD", "EURUSD", "EURUSD"],
    "freq": ["1H", "1H", "1H", "1H", "1H"],
    "open_price": [1.06969, 1.06930, 1.06929, 1.06932, 1.06952],
    "high_price": [1.06979, 1.06945, 1.06941, 1.06969, 1.06969],
    "low_price": [1.06930, 1.06930, 1.06929, 1.06932, 1.06949],
    "close_price": [1.06930, 1.06938, 1.06933, 1.06952, 1.06951],
    "volume": [12, 3, 31, 63, 70],
    "datetime": [
        pd.Timestamp("2023-01-02 07:00:00"),
        pd.Timestamp("2023-01-02 08:00:00"),
        pd.Timestamp("2023-01-02 09:00:00"),
        pd.Timestamp("2023-01-02 10:00:00"),
        pd.Timestamp("2023-01-02 11:00:00")
    ]
}

# 创建 DataFrame
df = pd.DataFrame(data).set_index("datetime")

# 确保 DataFrame 的索引具有明确的1小时频率
df = df.asfreq('1H')

# 移除非数值列
df = df.drop(columns=['symbol', 'freq'])

# 转换为 GluonTS 数据集
ds = PandasDataset(dict(df))

# 定义测试集长度
TEST = 100  # 根据你的数据量调整

# 划分训练集和测试集
train, test_template = split(ds, offset=-TEST)

# 定义预测长度和窗口数量
PDT = 20  # 预测长度
windows = TEST // PDT  # 窗口数量
distance = PDT  # 窗口间隔

# 生成滚动窗口评估数据
test_data = test_template.generate_instances(
    prediction_length=PDT,
    windows=windows,
    distance=distance
)

# 定义模型参数
SIZE = "base"  # 模型大小
CTX = 200  # 上下文长度
PSZ = "auto"  # 补丁大小
BSZ = 32  # 批量大小


# 准备预训练模型
model = MoiraiForecast(
    module=MoiraiModule.from_pretrained(
        f"Salesforce/moirai-moe-1.0-R-{SIZE}"),
    prediction_length=PDT,
    context_length=CTX,
    patch_size=PSZ,
    num_samples=100,
    target_dim=1,
    feat_dynamic_real_dim=ds.num_feat_dynamic_real,
    past_feat_dynamic_real_dim=ds.num_past_feat_dynamic_real,
)

# 创建预测器
predictor = model.create_predictor(batch_size=BSZ)

# 进行预测
with torch.no_grad():  # 确保在预测时不计算梯度
    forecasts = predictor.predict(test_data.input)

# 获取第一个输入、标签和预测结果
input_it = iter(test_data.input)
label_it = iter(test_data.label)
forecast_it = iter(forecasts)

inp = next(input_it)
label = next(label_it)
forecast = next(forecast_it)

# 绘制预测结果
fig, ax = plot_single(
    inp,
    label,
    forecast,
    context_length=CTX,
    name="pred",
    show_label=True
)

# 添加图例
ax.legend(['Input Data (History)', 'True Label (Actual)', 'Forecast (Predicted)'])

# 显示图形
plt.show()
